import math
import sys

from PyQt6.QtCore import Qt
from PyQt6.QtGui import QVector3D, QLinearGradient
from PyQt6.QtWidgets import QApplication, QWidget, QMainWindow
from PyQt6.QtDataVisualization import Q3DSurface, QSurface3DSeries, QSurfaceDataProxy, QSurfaceDataItem,\
                                      QValue3DAxis, Q3DCamera, Q3DTheme


class MyWidget(QMainWindow):
    def __init__(self):
        super(MyWidget, self).__init__()
        self.setWindowTitle('3D图表：三维空间电子衍射图像')

        # 创建图表
        surface = Q3DSurface()
        container = QWidget.createWindowContainer(surface)
        self.setCentralWidget(container)
        cameraView = Q3DCamera.CameraPreset.CameraPresetLeftHigh
        surface.scene().activeCamera().setCameraPreset(cameraView)

        # 封装数据
        proxy = QSurfaceDataProxy()
        x = -20
        for i in range(400):
            row = []
            z = -20
            for j in range(400):
                y = math.cos(math.sqrt(x * x + z * z)) * 2.5
                v3d = QVector3D(x, y, z)
                item = QSurfaceDataItem(v3d)
                row.append(item)
                z += 0.1
            x += 0.1
            proxy.addRow(row)
        series = QSurface3DSeries()
        series.setDataProxy(proxy)
        series.setDrawMode(QSurface3DSeries.DrawFlag.DrawSurface)

        # 图形颜色
        gradient = QLinearGradient()
        gradient.setColorAt(1.0, Qt.GlobalColor.yellow)
        gradient.setColorAt(0.5, Qt.GlobalColor.cyan)
        gradient.setColorAt(0.2, Qt.GlobalColor.red)
        gradient.setColorAt(0.0, Qt.GlobalColor.lightGray)
        series.setBaseGradient(gradient)
        series.setColorStyle(Q3DTheme.ColorStyle.ColorStyleRangeGradient)
        surface.addSeries(series)

        # 坐标轴
        axisX = QValue3DAxis()
        axisX.setTitle('X')
        axisX.setTitleVisible(True)
        axisX.setRange(-21, 21)
        surface.setAxisX(axisX)

        axisY = QValue3DAxis()
        axisY.setTitle('Y')
        axisY.setTitleVisible(True)
        axisY.setRange(-10, 10)
        surface.setAxisY(axisY)

        axisZ = QValue3DAxis()
        axisZ.setTitle('Z')
        axisZ.setTitleVisible(True)
        axisZ.setRange(-21, 21)
        surface.setAxisZ(axisZ)


if __name__ == '__main__':
    app = QApplication(sys.argv)
    ex = MyWidget()
    ex.show()
    sys.exit(app.exec())
